/* Copyright (c) 2011-2013 ISP RAS, Moscow. For more information please write to

      mermakov@ispras.ru
      svartanov@ispras.ru

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
*/

package ru.ispras.coffeemachine.core;

import java.io.BufferedReader;
import java.io.File;
import java.io.FileOutputStream;
import java.io.FileReader;
import java.io.FileWriter;
import java.io.IOException;
import java.io.PrintWriter;
import java.text.ParseException;
import java.util.ArrayList;
import java.util.List;

import org.kohsuke.args4j.CmdLineParser;

import ru.ispras.coffeemachine.instrument.Espresso;
import ru.ispras.coffeemachine.target.Concurrency;
import ru.ispras.coffeemachine.trace.ComplexQueries;
import ru.ispras.coffeemachine.trace.Glace;
import ru.ispras.coffeemachine.trace.PostProcessor;
import ru.ispras.coffeemachine.util.Executor;
import ru.ispras.coffeemachine.util.Logger;
import ru.ispras.coffeemachine.util.Options;
import ru.ispras.coffeemachine.util.Output;
import ru.ispras.coffeemachine.util.Util;

/**
 * Coffee Machine, instrument for dynamic analysis Java-programs
 *
 * @author Sergey Vartanov (me@enzet.ru)
 * @author Mikhail Ermakov (mermakov@ispras.ru)
 * @version 1.2.0
 * @since 2011-10-03
 */
public class CoffeeMachine
{
	public class Iteration
	{
		public List<Boolean> points;
		public int STPNumber;
		public int heuristic;

		public Iteration(List<Boolean> points, int STPNumber, int instructionsCount)
		{
			this.points = points;
			this.STPNumber = STPNumber;
			heuristic = instructionsCount;
		}

		@Override
		public String toString()
		{
			return "heuristic: " + heuristic;
		}
	}

	static int iterationNumber = 1;
	static int exploitNumber = 1;

	static List<Iteration> iterations;
	public static PrintWriter codePositions;

	private static String file;

	public static void main(String[] argv) throws Exception
	{
		try
		{
			new Logger();
			new Executor();

			try
			{
				CmdLineParser parser = new CmdLineParser(new Options());
				parser.parseArgument(argv);

				Options.check();

				if (Options.printHelp)
				{
					parser.printUsage(System.out);
					deleteTemp();
				}
			}
			catch (Exception e)
			{
				Logger.error(e.getMessage());
				Logger.info("Type -h or --help for help.");
				return;
			}
			if (!Options.printHelp)
			{
				machine();
			}
		}
		catch (Exception e)
		{
			e.printStackTrace();
			deleteTemp();
		}
	}

	/**
	 * Instrument target classes;
	 * Run first iteration of analysis unless <code>--instrument</code> option was specified
	 * @throws Exception
	 */
	public static void machine() throws Exception
	{
		Espresso espresso = new Espresso();

		// Clear instrumentation directory and copy rt.jar to the own directory

		boolean isInstrument = (Options.instrumentationID == -1) || Options.instrumentOnly;

		if (isInstrument)
		{
			Util.makeDir(Options.instrumentationDir);

			if (Options.rtPath != null)
			{
				Util.copyFile(Options.rtPath, Options.instrumentationDir + "rt.jar");
			}
			codePositions = new PrintWriter(new FileOutputStream(Options.codePositionFile, false), true);

			espresso.instrument();
			espresso.update();
		}
		if (Options.instrumentOnly)
		{
			codePositions.close();
			return;
		}
		if (Options.inputs.size() > 0) {
			Util.makeDir(Options.tempDir);

			for (int i = 0; i < Options.inputs.size(); i++) {
				int index = Options.inputs.get(0).lastIndexOf('/');
				String file = Options.inputs.get(i).substring(index + 1);
				Util.copyFile(Options.inputs.get(i), Options.tempDir + file);
				index = Options.arguments.indexOf(Options.inputs.get(i));
				if (index != -1) {
					Options.arguments.set(index, Options.tempDir + file);
				}
				Options.inputs.set(i, Options.tempDir + file);
			}
		}
		iterations = new ArrayList<Iteration>();

		iteration(0);

		deleteTemp();
	}

	/**
	 * Run analysis iteration:
	 * <ol>
	 * 		<li>Execute target application</li>
	 * 		<li>Check errors</li>
	 * 		<li>Parse STP output&mdash;generate new input files</li>
	 * 		<li>Recursively call next iteration</li>
	 * </ol>
	 */
	public static void iteration(int depth) throws Exception
	{
		Logger.info("\nIteration " + iterationNumber + ": " + depth);

		if (Options.inputs.size() > 0)
		{
			file = Executor.execute("cat " + Options.inputs.get(0)).getOut();

			String content = " ";
			for (int i = 0; i < file.length(); i++)
			{
				content += String.format("%4x", (int)file.charAt(i)) + " ";
			}
			Logger.info("File: " + content + ".");
		}
		// Application execution

		StringBuilder command = new StringBuilder();

		command.append("java -Xbootclasspath/p:" + Options.instrumentationDir + "rt.jar -cp " + Options.instrumentationDir);

		for (String jar : Options.classfiles.getJars())
		{
			command.append(":").append(Options.instrumentationDir).append(jar);
		}

		if (Options.targetClassfile == null)
		{
			command.append(" " + Options.targetJar);
		}
		else
		{
			command.append(" " + Options.targetClassfile);
		}

		for (String argument : Options.arguments)
		{
			command.append(" ").append(argument);
		}
		System.out.println("Command: " + command);

		Logger.verbose("Execution: " + command.toString() + ".");

		Output output = Executor.execute(command.toString());

		System.out.println(output.getOut());

		if (Options.eventsFileName != null && Options.concurrencyEnabled) {
			try {
				File eventsFile = new File(Options.eventsFileName);
				FileWriter eventsFileWriter = new FileWriter(eventsFile);

				String trace = output.getFilteredOutput(Concurrency.PREFIX);

				boolean isInstrument = (Options.instrumentationID == -1) || Options.instrumentOnly;

				StringBuffer cps = new StringBuffer();

				if (!isInstrument)
				{
					BufferedReader reader = new BufferedReader(new FileReader(Options.codePositionFile));

					String line = "";
					while ((line = reader.readLine()) != null)
					{
						cps.append(line + "\n");
					}
					reader.close();
				}
				trace = new PostProcessor(cps.toString(), trace).postprocess();

				eventsFileWriter.append(trace);
				eventsFileWriter.close();
			} catch (IOException e) {
				System.err.println("Cannot write TSan event file");
			}
		}

		if (!output.getErr().equals(""))
		{
			Logger.exploit(output.getErr(), exploitNumber);
			dumpExploit(output.getErr());

			System.out.println(output.getErr());
		}

		//Util.write("../../trace", output.getOut());

		//System.out.println(output.getOut());

		ComplexQueries complexQueries = new ComplexQueries(output.getOut());
		complexQueries.parse(depth);

		//Util.write(Options.eventsFileName, Concurrency.parse(output.getOut()));

		if (Options.doSingleRun)
		{
			return;
		}
		System.out.println("Next branches: " + complexQueries.getSize());

		while (complexQueries.hasQuery())
		{
			String query = complexQueries.getQuery();
			int newDepth = complexQueries.getDepth();

			Util.write(Options.tempDir + "/queries", query);

			//Logger.log(query);

			Output trace;
			trace = Executor.execute(Options.stp + " -p " + Options.tempDir + "/queries");

			Glace glace = new Glace(trace.getOut(), 1000);

			boolean valid = false;

			if (Options.printQueries)
			{
				Logger.info("Full query:\n" + query);
			}

			try
			{
				valid = glace.parse();
			}
			catch (ParseException e)
			{
				Logger.error("parsing STP answer: " + e);

				if (!Options.printQueries)
				{
					Logger.info("Full query:\n" + query);
				}
				Logger.info("Full answer:\n" + trace.getOut() + trace.getErr());
				deleteTemp();
			}
			catch (Exception e)
			{
				Logger.error(e.toString());
			}

			if (!valid)
			{
				if (!Options.inputs.isEmpty())
				{
					Util.write(Options.inputs.get(0), glace.files.get(0));
				}

				iterationNumber++;

				iteration(newDepth);
			}
		}
	}

	/**
	 * Exploit dumping
	 */
	private static void dumpExploit(String error) throws IOException, InterruptedException
	{
		Output trace;

		Executor.execute("mkdir exploits");
		Executor.execute("mkdir exploits/" + exploitNumber);

		if (!Options.inputs.isEmpty())
		{
			trace = Executor.execute("cp " + Options.inputs.get(0) + " exploits/" + exploitNumber + "/" + Executor.trimDirectory(Options.inputs.get(0)));
			trace.checkErrors();
		}

		Util.write("exploits/" + exploitNumber + "/exploit", error);

		exploitNumber++;
	}

	public int getExploits()
	{
		return exploitNumber - 1;
	}

	public int getIterations()
	{
		return iterationNumber;
	}

	/**
	 * Delete temporary files
	 */
	public static void deleteTemp()
	{
		if (!Options.keepTemp && Options.instrumentationID == -1)
		{
			Util.deleteFileRecursive(Options.instrumentationDir);
			Util.deleteFileRecursive(Options.tempDir);
		}
	}
}
